from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from major.models_manager import chat_model
from graphs import make_png
llm = chat_model.get_model().bind(max_tokens=32,temperature=2)
# 状态
class State(TypedDict):
    topic: str
    joke: str
    improved_joke: str
    generate_counts: int

# 节点
def generate_joke(state: State):
    msg = llm.invoke(f"请写一个关于{state['topic']}的简短中文笑话")
    return {"joke": msg.content, "generate_counts":state.get("generate_counts", 0)+1}

def improve_joke(state: State):

    msg = llm.invoke(f"为这个笑话添加一个意想不到的转折: {state['joke']}, 仅输出提高后的笑话")
    return {"improved_joke": msg.content}

def check_joke(state: State):
    if state["generate_counts"] > 2:
        return "Pass"

    # 简单模拟笑点检查
    if "？" in state["joke"]:
        return "Pass"
    else:
        return "Fail"

# 1.初始化图
workflow = StateGraph(State)
# 2.添加节点
workflow.add_node("generate_joke", generate_joke) # 传节点名字，节点执行方法
workflow.add_node("improve_joke", improve_joke)
# 3.添加连接边
workflow.add_edge(START, "generate_joke")
workflow.add_edge("generate_joke", "improve_joke")
workflow.add_conditional_edges("improve_joke", # 上一个节点的名字
                               check_joke, # 条件方法
                               {"Pass": END, "Fail": "generate_joke"} )

# 4.编译图
graph = workflow.compile()

# 5.保存(可选)
make_png.save(graph, "有环图.png")
# 执行
state = graph.invoke({"topic": "猫"},print_mode='updates')

print(state['improved_joke'])
